package com.wutong.erb.filter;

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import lombok.Getter;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Component;
import org.springframework.web.filter.OncePerRequestFilter;

import java.io.IOException;
import java.time.Duration;

/**
 * 限流过滤器
 * @author wutong
 */
@Getter
@Component
public class RateLimitFilter extends OncePerRequestFilter {
    private final Bucket bucket;
    private final String[] WHITE_LIST = {
            "/api/auth/check-username",
            "/api/auth/check-email"
    };

    public RateLimitFilter() {
        Bandwidth limit = Bandwidth.builder()
                // 令牌桶容量
                .capacity(40)
                // 每分钟补充5个令牌
                .refillIntervally(20, Duration.ofMinutes(1))
                .build();
        this.bucket = Bucket.builder()
                .addLimit(limit)
                .build();
    }
    @Override
    protected void doFilterInternal(@NotNull HttpServletRequest request,
                                    @NotNull HttpServletResponse response,
                                    @NotNull FilterChain chain) throws IOException, ServletException {
        if (shouldFilter(request)) {
            chain.doFilter(request, response);
            return;
        }
        if (bucket.tryConsume(1)) {
            chain.doFilter(request, response);
        } else {
            response.setHeader("X-RateLimit-Retry-After", "60");
            response.sendError(429, "请求太频繁");
        }
    }

    private boolean shouldFilter(HttpServletRequest request) {
        for (String white : WHITE_LIST) {
            if (request.getRequestURI().startsWith(white)) {
                return false;
            }
        }
        return true;
    }
}


